import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.linear_model import LinearRegression
from sklearn.metrics import mean_squared_error
import numpy as np
import matplotlib.pyplot as plt

def getTrainSetAndTestSet(DataPath): 
    data = pd.read_csv(DataPath) 
    X = data[['AT','V','AP','RH']] # AT, V, AP 和 RH 这 4 个列作为样本特征 
    y = data[['PE']] # 用 PE 作为样本输出 
    X_train, X_test, y_train, y_test = train_test_split(X, y, random_state=1) # 随机划分训练集和测试集，默认 25% 作为测试集 
    print("训练集维度：", X_train.shape, y_train.shape)  # 查看训练集和测试集的维度
    print("测试集维度：", X_test.shape, y_test.shape)
    return X_train, X_test, y_train, y_test

def TrainLinearRegression(X_train, y_train): 
    linreg = LinearRegression() # 未经训练的机器学习模型 
    linreg.fit(X_train, y_train) # 训练模型
    print("截距：", linreg.intercept_)  # 输出线性回归的截距
    print("系数：", linreg.coef_)  # 输出各个系数
    return linreg 

def EvaluationModel(linreg, X_test, y_test): 
    y_pred = linreg.predict(X_test) 
    print("均方误差 MSE：", mean_squared_error(y_test, y_pred))  # 输出均方误差 MSE
    print("均方根误差 RMSE：", np.sqrt(mean_squared_error(y_test, y_pred)))  # 输出均方根误差 RMSE
    return y_pred 

def Visualization(y_test, y_pred): 
    fig, ax = plt.subplots() 
    ax.scatter(y_test, y_pred) 
    ax.plot([y_test.min(), y_test.max()], [y_test.min(), y_test.max()], 'k--', lw=5) 
    ax.set_xlabel("Measured") 
    ax.set_ylabel("Predicted") 
    plt.show()

# 主程序
if __name__ == "__main__":
    # 指定数据路径（D盘根目录下的 Folds5x2_pp.csv）
    data_path = r"D:\Folds5x2_pp.csv"  # 使用原始字符串（r）避免转义问题
    
    # 获取训练集和测试集
    X_train, X_test, y_train, y_test = getTrainSetAndTestSet(data_path)
    
    # 训练线性回归模型
    linreg = TrainLinearRegression(X_train, y_train)
    
    # 评估模型
    y_pred = EvaluationModel(linreg, X_test, y_test)
    
    # 可视化预测结果
    Visualization(y_test, y_pred)